[WS1][CUDA][Qwen3-Next] Gated RMSNorm CUDA kernel for the GDN block (RFC #428 stack 3/4) - #468

tests/test_rms_norm.py                         57 passed, 72 skipped in 14.12s
tests/test_qwen3_next_norm.py                  141 passed in 9.03s
tests/test_qwen3_next_workload.py              7 passed in 4.13s
tests/test_ws1_gtest_gpu.py                    9 passed in 54.17s
tests/test_vjp_fp32.py                         50 passed in 5.58s
rl_engine/tests/test_dispatch.py               17 passed in 5.87s
tests/check_qwen3_next_norm_providers.py       9 passed, 14 warnings in 62.98s (0:01:02)
rms_norm               suite=rms_norm passed=True pass_rate=1.0000
rms_norm_gated         suite=rms_norm_gated passed=True pass_rate=1.0000
qwen3_next_rms_norm    suite=qwen3_next_rms_norm passed=True pass_rate=1.0000
op=qwen3_next_rms_norm profile=cuda_bf16 candidate=rl_engine.kernels.ops.cuda.norm.rmsnorm.Qwen3NextRMSNormCudaOp::cuda passed=True
op=qwen3_next_rms_norm profile=cuda_bf16 candidate=rl_engine.kernels.ops.cuda.norm.rmsnorm.Qwen3NextRMSNormCudaOp::cuda passed=True
op=rms_norm_gated profile=cuda_bf16 candidate=rl_engine.kernels.ops.cuda.norm.rmsnorm.Qwen3NextRMSNormGatedCudaOp::cuda passed=True
op=rms_norm_gated profile=cuda_bf16 candidate=rl_engine.kernels.ops.cuda.norm.rmsnorm.Qwen3NextRMSNormGatedCudaOp::cuda passed=True

Each gate prints one accuracy, invariance or singleton_aggregate line per comparison; the four logs total 72 invariance and singleton-aggregate comparisons, all max_abs=0.

The rstd identity test is meant to fail if anything perturbs the gated statistic. A separate run
checked that at 5eaeca8: it added +1e-4f to each square in the gated kernel's
sum of squares and rebuilt. The rstd test then failed in all 36 configurations
(36 failed, 100 deselected), so the test detects a perturbed gated statistic.

The edited line, at 5eaeca8 (git show 4ddbaa0:csrc/cuda/rmsnorm.cu | sed -n 245,256p):

251:        local_sum += xv * xv;        // inside rmsnorm_gated_fwd_kernel

At b0f4c5c the same line is :254 (the PR below added a three-line comment at
:146-148; the plain kernel's copy at :125 is untouched). To repeat on this branch:

sed -i '254s/local_sum += xv \* xv;/local_sum += xv * xv + 1e-4f;/' csrc/cuda/rmsnorm.cu
git diff --stat   # one line changed in csrc/cuda/rmsnorm.cu
MAX_JOBS=32 TORCH_CUDA_ARCH_LIST=10.0 RL_KERNEL_REQUIRE_EXT=1 \
    pip install --no-build-isolation --no-deps -e .
python -m pytest tests/test_qwen3_next_norm.py -q -p no:randomly -k rstd_is_bitwise
git checkout csrc/cuda/rmsnorm.cu   # then rebuild

At 5eaeca8 the file collected 136 tests, hence 100 deselected; at b0f4c5c it collects
141. The recipe has not been re-run at b0f4c5c.

Two Python-only mutations (no rebuild), each run on its own against the full
tests/test_qwen3_next_norm.py on a B200 with _C loaded from the tree under test:

Each new test fails when the behavior it pins is removed, and no other test in the file
catches either change. Restore with git checkout rl_engine/kernels/ops/cuda/norm/rmsnorm.py.

Cluster, 5eaeca8, same software as above. Each run built _C from its own clean
snapshot of the ref and refused to run unless rl_engine._C.__file__ was inside that
snapshot.

• Full suite: 15 failed, 3098 passed, 923 skipped; no new failures
  against main, and the two fixed by the bottom PR stay fixed. Arithmetic: 2942 + 2
  fixed + 154 newly collected (4036 − 3882) = 3098; 83 of the 154 are added in this PR.
• Targeted pytest at 5eaeca8: test_qwen3_next_norm.py 136 passed (71 new here,
  incl. the 36 rstd cases); test_dispatch.py 17 passed (4 new);
  test_qwen3_next_workload.py 7 passed; check_qwen3_next_norm_providers.py 9 passed;
  test_rms_norm.py 57 passed, 72 skipped; test_ws1_gtest_gpu.py 9 passed;
  test_vjp_fp32.py 50 passed.
• C3/C4 gates: all four report passed=True. All 72 comparisons are
  bitwise (max_abs=0): 68 invariance pairs across batch size, chunking, padding and
  permutation, plus 4 singleton-aggregate checks of dweight, which supply dweight's
  batch-size coverage. This was the first time these gates ran anywhere; CI has not
  reached them.
• Mutation check: above.
• check_operator, without shape flags (the defaults above): rms_norm_gated,
  qwen3_next_rms_norm and rms_norm, each passed=True pass_rate=1.0000. The
  job is not recorded.

Local, aaf0986: 2× B200, same software; extension built with
-gencode=arch=compute_100,code=sm_100. The _C it loaded was built from 82e6c16
(git diff --stat de58159 82e6c16 -- csrc/ is empty); that build has since been deleted,
so this cannot be re-checked. The two commits after aaf0986 and before the merge
change only tests.

full suite, local main (32b765e): 19 failed, 2940 passed, 923 skipped
full suite, local de58159:        17 failed, 3059 passed, 923 skipped   (no new failures)

main has two fewer failures on the cluster node than locally; each comparison is
against the same-environment baseline. Measured at aaf0986 against main: real
pre-commit, no new findings (also clean on the two later test commits);
mypy --ignore-missing-imports rl_engine/, no new errors (not re-run since);
mkdocs build --strict, 8 warnings, the same 8 as main.

Qwen3NextRMSNormGatedCudaOp.__init__ checks for the compiled symbols
(rmsnorm.py:325-327). The registry's _load_backend (registry.py:1372) only catches
import and attribute errors, but _get_or_create_backend (:1353-1370) skips a backend
whose construction raises. Raising in __init__ therefore lets the CUDA-first list fall
through to the PyTorch reference on a build without the extension. Neither runtime hook
(_adjust_priority_from_env, _adjust_priority_for_hardware) inserts another backend
for these two ops.

rl_engine/tests/test_dispatch.py:284-325 covers both ops. Since b0f4c5c the
fallback test resolves for device="cuda", so the CUDA-first list is walked even on a
CPU-only host, and it asserts the CUDA backend is in _failed_backends. It is skipped on
HIP builds.

Both specs carry a cuda-sm90 candidate, as rms_norm, qk_norm and lm_head do
(det_gemm, also a reduction spec, does not). Without it, --candidate cuda-sm90 fails
with a KeyError.

• csrc/ops.cpp (7dd291a; duplicates removed in b0f4c5c), checks shared by the
  plain and gated entry points (:313-334):
  • x must be fp32, fp16 or bf16 (fp64 now rejected by rmsnorm_forward and
    rmsnorm_backward_dx); weight must have the x dtype or be fp32.
  • Backward: dy must have the x dtype, and rstd must be fp32.
    rmsnorm_backward_dw has these checks plus device, mask and H checks, but no x
    dtype restriction.
  • All: tensors on the same device (including mask), and H > 0.
  • The launchers that lacked them gained a CUDAGuard and
    C10_CUDA_KERNEL_LAUNCH_CHECK; the left-fold launcher already had them. The per-row
    launchers and the chunked partial_dw launcher return early when T == 0.
  • b0f4c5c drops the dim/size checks rmsnorm_check_weight already performs; error
    messages for those cases now come from it, and no test matches the removed messages.
• USE_ROCM guard (7dd291a, 0cceb84). setup.py compiles rmsnorm.cu for ROCm
  as well. Static check only, no ROCm device: PyTorch's hipify (extension mode)
  translates the file without error; the two headers and C10_CUDA_KERNEL_LAUNCH_CHECK
  are already used unguarded by sources the ROCm build compiles (activation.cu,
  deterministic_attention.cu, the logp kernels); and since 0cceb84 the launchers use
  the same at::cuda::OptionalCUDAGuard as activation.cu. The only c10::cuda::CUDAGuard
  left is inside the USE_ROCM-excluded left-fold block. A ROCm build has not been run.
• Signed zero. The gated kernels add weight_offset only when it is nonzero
  (aaf0986), as the plain kernels do after the PR below. Since b0f4c5c the Python
  dgate path does the same (rmsnorm.py:270-273), pinned by
  test_cuda_gated_dgate_keeps_signed_zero_weight (3 dtypes).
• Shared gate scripts (8df1ffb, b0f4c5c). scripts/check_{forward,gradient}_invariance.py
  gain --manifest, and their JSON report gains a workload block for every op.
  validate_norm_dimensions returns immediately for any manifest other than the
  Qwen3-Next one (qwen3_next_workload.py:87-93). Since b0f4c5c,
  workload.full_model_evidence is read from the manifest: False for the Qwen3-Next
  manifest, None for the Dense manifest, which declares none (it was hard-coded
  False for every op).
• rl_engine/testing/ws1_workload.py: load_manifest dispatches validation on the
  manifest's scope, and _validate_model_identity takes the fingerprint as a
  parameter (:315-320).
• Shared adapter semantics: adapters can be model-scoped. An existing Ascend test
  (tests/test_ws1_ascend_closeout.py:195-199) skips adapters scoped to another model,
  and a new test (:207) asserts those resolve to absent_not_required.

rl_engine/testing/qwen3_next_norm_manifest.json pins Qwen/Qwen3-Next-80B-A3B-Instruct
at revision 9c7f2fbe…, with 41 shard SHA-256s and sizes, the index file's SHA-256, and a
content_hash. How these values were first obtained is not recorded on this branch
(8df1ffb's message is subject-only), and nothing on the branch re-derives them from
files. What the validator checks (ws1_workload.py:337-369, called from
validate_norm_manifest, qwen3_next_workload.py:44) is internal consistency only:

• every shard digest and index_sha256 is 64 lowercase hex characters;
• weight_files_total_size_bytes equals the sum of the shard sizes;
• content_hash equals weight_snapshot_hash over the sorted (filename, sha256, size)
  records (:1135-1139).

The values are, however, reproducible from public Hugging Face Hub metadata at that
revision: all 41 shard SHA-256s and sizes equal the LFS oid/size of the revision's 41
.safetensors files; index_sha256 is the SHA-256 of its model.safetensors.index.json;
total_size_bytes is that index's metadata.total_size; and content_hash recomputes.
No weight bytes were downloaded, so this relies on Hugging Face's LFS digests.

REV=9c7f2fbe84465e40164a94cc16cd30b6999b0cc7
curl -sSL "https://huggingface.co/api/models/Qwen/Qwen3-Next-80B-A3B-Instruct/tree/$REV" -o tree.json
curl -sSL "https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct/resolve/$REV/model.safetensors.index.json" | sha256sum
python - <<'PY'
import json
from rl_engine.testing.ws1_workload import weight_snapshot_hash
hf = {e["path"]: e.get("lfs") or {} for e in json.load(open("tree.json"))}
ws = json.load(open("rl_engine/testing/qwen3_next_norm_manifest.json"))["model_identity"]["weight_snapshot"]
print(sum(hf[s["filename"]].get("oid") == s["sha256"] and hf[s["filename"]].get("size") == s["size_bytes"]
          for s in ws["shards"]), "/", len(ws["shards"]))
print(weight_snapshot_hash(ws["shards"]) == ws["content_hash"])
PY

The gates never load weights: they use synthetic inputs at the checkpoint's dimensions
(provenance_boundary.runtime_verified: false, full_model_evidence: false). The Dense
path's verify_hf_weight_snapshot (rl_engine/alignment/qwen3_dense.py:548) has no
Qwen3-Next counterpart.

vLLM has three numerical paths for this norm, and none is bitwise equal to another:
A. eager forward_native (PyTorch ops); B. FLA forward_cuda (Triton); C.
forward_native compiled by Inductor inside the model graph.

• Default (inductor) mode → C. custom_ops defaults to "none"
  (config/vllm.py:1608-1615). The norm is called at qwen_gdn_linear_attn.py:856,
  outside the splitting op vllm::qwen_gdn_attention_core, so it is traced into the
  model graph.
• enforce_eager → B. custom_ops becomes "all" (config/vllm.py:1615).

A vs B (committed in tests/check_qwen3_next_norm_providers.py, from the PR below and
not changed here): over 40 seeds (bf16, head_v_dim=128, 512 rows, one shape, B200)
they disagreed on 21, worst 1.56e-2. In fp32 they differ on about 36% of elements,
under an asserted bound of 1e-5. The file bounds the eager gaps rather than asserting
equality.

Conventions. vLLM's fp32 weight multiply is Qwen3NextRMSNormGatedOp;
transformers' cast-first multiply is Qwen3NextRMSNormGatedHFOp. One earlier
measurement, with no committed script, found them differing on 35% of elements in bf16,
max 6.25e-2 (docs/operators/qwen3-next-rms-norm-gated.md:91-94, one seed). The
committed test asserts only that they differ (> 1e-3) in bf16 and agree bitwise in
fp32.

Question (RFC #428 §0 item 1): which source of truth should the CUDA strict profile
use for forward arithmetic?

• The RFC's §0 profile table and §7.2 ("GDN gated RMSNorm | Shared forward") point to a
  shared vLLM forward. For this norm that means the FLA kernel or the compiled path,
  depending on mode.
• The existing Qwen3 dense strict profile instead patches vLLM's RMSNorm to PyTorch
  eager arithmetic on CUDA (rl_engine/integrations/vllm_runtime.py:1283-1304,
  :1415); on ROCm strict_rms_norm goes to rocm.rmsnorm instead
  (rl_engine/kernels/ops/pytorch/norm/rms_norm.py:17-21, and :44-48 with a residual).

The answer also decides which gated convention is the strict default.

The Triton op still has the split the plain CUDA op had on main: its parameter-VJP
hook recomputes rstd with mean(-1) (rl_engine/kernels/ops/triton/rmsnorm_triton.py:164-169)
while its backward uses the saved rstd (:127-130). That is how it is on main, and
it is left for a separate change.

From a one-off standalone probe (not committed, not an engine run):

• C vs A: 20/40 seeds differ;
• the gated PyTorch reference vs C: 17–18/40;
• the gated PyTorch reference vs A: 6/40.

On VLLM_BATCH_INVARIANT=1:

• It does not reach this norm in either mode. Under enforce_eager the FLA kernel never
  calls aten.mean.
• In the standalone compile, output was bitwise identical with the flag on and off
  (40/40), with 0 calls to the batch-invariant mean.
• Calling forward_native directly in eager mode with the flag on also left the output
  bitwise unchanged (0/40).

The same probe found no prefix-slice batch-invariance failure for vLLM's eager path or
its dim-0-dynamic compile (seeds 0–7, n up to 256). See the C1 PR for the full table and
its scope limits.

As of 2026-10-01 no CI path executes these tests, before or after merge.

• On main (32b765e and the current 1968a87), CI-Pipeline's linting fails at
  Run pre-commit hooks, so unit-tests is skipped. Docs fails at both commits too.
• WS1-chain-GPU and ws1-gtest-gpu stop at Configure runpodctl. The cause is step
  order, not a recent release: runpodctl config --apiKey, run as the first runpodctl
  call on a fresh runner, finds no config file (run 36835086246 logs
  'runpodctl config' is deprecated then error saving config, reproduced locally with
  v2.14.0 on an empty HOME). ws1-gtest-gpu has never succeeded: as of 2026-10-01,
  122 runs (69 failed, 40 skipped, 13 action_required), and its first run (2026-08-19)
  already failed at this step. A separate fix is being prepared.
• gpu-ci (pull_request_target only, gpu-ci.yml:4; label needs-gpu-ci) uses the
  same command (:50) but runs runpodctl version first, which creates the config
  file, so it got past setup on its last run that reached
  RunPod (2026-09-11, run 34585491929) and then failed at pod create. No run since, so
  whether it still gets past setup is unverified. Its pytest tests/
  (ci/run_gpu_ci.sh:133) would also stop at the aiter collection error.
• ws1-chain-npu was queued on 1968a87 and cancelled on 32b765e.

If those are fixed, the workflows' scripts would run:

1. ci.yml (needs: linting): tests/test_qwen3_next_norm.py and
   rl_engine/tests/test_dispatch.py on a CPU-only runner, so the CUDA cases skip; the
   dispatch fallback test does walk the CUDA list there. Not test_qwen3_next_workload.py.
2. ws1-gtest-gpu.yml: same-repo PRs and push to main, skipped for fork PRs
   (:59); this branch touches its paths. It would run the four Qwen3-Next gates
   (ci/run_ws1_gtest.sh:79-88) and tests/test_ws1_gtest_gpu.py, none of the other
   pytest files above. It runs on sm_86 (RTX A4000) and sm_90 (H100), not sm_100, under
   set -euo pipefail with the Qwen3-Next stage last, so any earlier failure skips it.
3. ws1-chain-gpu.yml: push and PR to main/test, no paths filter, skipped for
   fork PRs (:46); maintainers can dispatch it on a fork's commit (:14-24). Its chain
   constructs RMSNormCudaOp directly (rl_engine/alignment/qwen3_dense.py:374-396), so
   it would execute the plain-kernel input checks this PR adds in csrc/ops.cpp.
4. gpu-ci.yml: with the collection error fixed, it would collect the CUDA cases but
   not the check_ file.

• sm_100 (Blackwell), not sm_90. Per RFC §2.2 no claim is made across H100/B200.
• Stacked on feat/cuda-qwen3-next-c1-norm, which is on
  fix/registry-cuda-rmsnorm-dispatch; both must merge first, and this branch contains
  their commits. cf68fd9 merges cleanly into the current upstream/test-qwennext (11cac8c,
  git merge-tree --write-tree). One further PR builds on this one.

• Duplicated sum-of-squares loop. The gated forward copies the plain kernel's
  sum-of-squares loop, block_reduce_sum and rsqrtf (rmsnorm.cu:248-259 vs
  :120-131), and the gated dx copies the dot loop (:302-309 vs :176-182). The
  bitwise rstd guarantee therefore rests on the 36-case test rather than on shared
  code. A shared device function is a follow-up.
• Private validator imports. qwen3_next_workload.py:9-18 imports six
  underscore-private validators from ws1_workload, and ws1_workload.py:268 imports
  back inside load_manifest to avoid the import cycle.
• Hard-coded dimensions. validate_norm_dimensions checks hidden != 2048 or head_dim != 128 (qwen3_next_workload.py:90), values already present in NORM_OPS
  (:37) and FINGERPRINT["linear_value_head_dim"] (:31).
• 5eaeca8's and b0f4c5c's messages say their CUDA tests had not yet run on a GPU;
  the results above supersede them. History on this branch is not rewritten.

None of these changes behavior. They are left for follow-ups, because any commit here
would have to be merged forward into the c6 branch stacked on it.